package main

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/valyala/fasthttp" // 配置常量
)

const (
	localIssuer   = "http://localhost:8020"
	clientID      = "demo-client"
	clientSecret  = "demo-secret"
	redirectURI   = "http://localhost:8020/callback"
	listenAddress = ":8020"
)

// RSA密钥对
var (
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
)

// 用户会话存储
var (
	sessions = make(map[string]*UserSession)
	users    = map[string]*User{
		"admin": {
			ID:       "user-001",
			Username: "admin",
			Password: "admin123",
			Email:    "admin@example.com",
			Name:     "管理员",
		},
		"test": {
			ID:       "user-002",
			Username: "test",
			Password: "test123",
			Email:    "test@example.com",
			Name:     "测试用户",
		},
	}
)

// 数据模型
type (
	User struct {
		ID       string `json:"sub"`
		Username string `json:"preferred_username"`
		Email    string `json:"email"`
		Name     string `json:"name"`
		Password string `json:"-"`
	}

	UserSession struct {
		UserID   string    `json:"user_id"`
		Username string    `json:"username"`
		Email    string    `json:"email"`
		Name     string    `json:"name"`
		Expiry   time.Time `json:"exp"`
	}

	OIDCConfig struct {
		Issuer                string   `json:"issuer"`
		AuthorizationEndpoint string   `json:"authorization_endpoint"`
		TokenEndpoint         string   `json:"token_endpoint"`
		UserInfoEndpoint      string   `json:"userinfo_endpoint"`
		JWKSUri               string   `json:"jwks_uri"`
		ResponseTypes         []string `json:"response_types_supported"`
		SubjectTypes          []string `json:"subject_types_supported"`
		IDTokenAlgs           []string `json:"id_token_signing_alg_values_supported"`
		Scopes                []string `json:"scopes_supported"`
		Claims                []string `json:"claims_supported"`
	}

	AuthRequest struct {
		ResponseType string `json:"response_type"`
		ClientID     string `json:"client_id"`
		RedirectURI  string `json:"redirect_uri"`
		Scope        string `json:"scope"`
		State        string `json:"state"`
		Nonce        string `json:"nonce"`
	}

	TokenResponse struct {
		AccessToken string `json:"access_token"`
		TokenType   string `json:"token_type"`
		ExpiresIn   int    `json:"expires_in"`
		IDToken     string `json:"id_token"`
	}

	UserInfoResponse struct {
		Sub               string `json:"sub"`
		Name              string `json:"name"`
		Email             string `json:"email"`
		PreferredUsername string `json:"preferred_username"`
	}
)

var clientID_demo = "demo-client"
var clientSecret_demo = "demo-secret"

func main() {
	// 初始化RSA密钥
	if err := initRSAKeys(); err != nil {
		log.Fatalf("初始化RSA密钥失败: %v", err)
	}

	// 设置路由
	mux := func(ctx *fasthttp.RequestCtx) {
		path := string(ctx.Path())
		method := string(ctx.Method())

		switch {
		case path == "/" && method == "GET":
			rootHandler(ctx)
		case path == "/.well-known/openid-configuration" && method == "GET":
			discoveryHandler(ctx)
		case path == "/.well-known/jwks.json" && method == "GET":
			jwksHandler(ctx)
		case path == "/authorize" && method == "GET":
			authorizeHandler(ctx)
		case path == "/token" && (method == "POST" || method == "GET"):
			tokenHandler(ctx)
		case path == "/userinfo" && method == "GET":
			userInfoHandler(ctx)
		case path == "/login" && method == "GET":
			loginFormHandler(ctx)
		case path == "/callback" && method == "POST":
			callbackHandler(ctx)
		default:
			ctx.NotFound()
		}
	}

	log.Printf("🚀 本地OIDC服务启动于: %s", localIssuer)
	log.Printf("📋 可用端点:")
	log.Printf("   - OIDC配置: %s/.well-known/openid-configuration", localIssuer)
	log.Printf("   - JWKS文档: %s/.well-known/jwks.json", localIssuer)
	log.Printf("   - 授权端点: %s/authorize", localIssuer)
	log.Printf("   - Token端点: %s/token", localIssuer)
	log.Printf("   - UserInfo端点: %s/userinfo", localIssuer)
	log.Printf("   - 测试登录: %s/login", localIssuer)
	log.Printf("   - 回调地址: %s/callback", localIssuer)

	log.Fatal(fasthttp.ListenAndServe(listenAddress, mux))
}

// OIDC发现端点
func discoveryHandler(ctx *fasthttp.RequestCtx) {
	config := OIDCConfig{
		Issuer:                localIssuer,
		AuthorizationEndpoint: localIssuer + "/authorize",
		TokenEndpoint:         localIssuer + "/token",
		UserInfoEndpoint:      localIssuer + "/userinfo",
		JWKSUri:               localIssuer + "/.well-known/jwks.json",
		ResponseTypes:         []string{"code", "id_token", "token"},
		SubjectTypes:          []string{"public"},
		IDTokenAlgs:           []string{"RS256"},
		Scopes:                []string{"openid", "profile", "email"},
		Claims:                []string{"sub", "name", "email", "preferred_username"},
	}

	ctx.SetContentType("application/json")
	json.NewEncoder(ctx).Encode(config)
}

// 授权端点 - 显示登录表单
func authorizeHandler(ctx *fasthttp.RequestCtx) {
	clientID := string(ctx.QueryArgs().Peek("client_id"))
	redirectURI := string(ctx.QueryArgs().Peek("redirect_uri"))
	scope := string(ctx.QueryArgs().Peek("scope"))
	state := string(ctx.QueryArgs().Peek("state"))

	if clientID != clientID_demo {
		ctx.Error("Invalid client_id", fasthttp.StatusBadRequest)
		return
	}

	// 重定向到登录页面
	loginURL := fmt.Sprintf("%s/login?redirect_uri=%s&scope=%s&state=%s",
		localIssuer, redirectURI, scope, state)
	ctx.Redirect(loginURL, fasthttp.StatusFound)
}

// 登录表单
func loginFormHandler(ctx *fasthttp.RequestCtx) {
	redirectURI := string(ctx.QueryArgs().Peek("redirect_uri"))
	state := string(ctx.QueryArgs().Peek("state"))

	ctx.SetContentType("text/html; charset=utf-8")
	fmt.Fprintf(ctx, `<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>本地OIDC服务 - 登录</title>
    <style>
        body { font-family: Arial, sans-serif; max-width: 400px; margin: 50px auto; padding: 20px; }
        .form-group { margin-bottom: 15px; }
        label { display: block; margin-bottom: 5px; }
        input[type="text"], input[type="password"] { width: 100%%; padding: 8px; border: 1px solid #ddd; }
        button { width: 100%%; padding: 10px; background: #007bff; color: white; border: none; cursor: pointer; }
        button:hover { background: #0056b3; }
        .info { background: #f8f9fa; padding: 15px; margin-bottom: 20px; border-radius: 5px; }
    </style>
</head>
<body>
    <div class="info">
        <h2>🔐 本地OIDC服务</h2>
        <p><strong>测试用户：</strong></p>
        <ul>
            <li>用户名：admin / 密码：admin123</li>
            <li>用户名：test / 密码：test123</li>
        </ul>
    </div>
    
    <form method="POST" action="/callback">
        <input type="hidden" name="redirect_uri" value="%s">
        <input type="hidden" name="state" value="%s">
        
        <div class="form-group">
            <label>用户名：</label>
            <input type="text" name="username" required>
        </div>
        
        <div class="form-group">
            <label>密码：</label>
            <input type="password" name="password" required>
        </div>
        
        <button type="submit">登录</button>
    </form>
</body>
</html>`, redirectURI, state)
}

// 回调处理 - 验证用户并颁发令牌
func callbackHandler(ctx *fasthttp.RequestCtx) {
	var username string
	var password string
	var redirectURI string
	var state string
	if ctx.FormValue("username") != nil {
		username = string(ctx.FormValue("username"))
	}
	if ctx.FormValue("password") != nil {
		password = string(ctx.FormValue("password"))
	}
	if ctx.FormValue("redirect_uri") != nil {
		redirectURI = string(ctx.FormValue("redirect_uri"))
	}
	if ctx.FormValue("state") != nil {
		state = string(ctx.FormValue("state"))
	}

	// 验证用户
	user, exists := users[username]
	if !exists || user.Password != password {
		ctx.Error("Invalid credentials", fasthttp.StatusUnauthorized)
		return
	}

	// 创建会话
	sessionID := generateSessionID()
	sessions[sessionID] = &UserSession{
		UserID:   user.ID,
		Username: user.Username,
		Email:    user.Email,
		Name:     user.Name,
		Expiry:   time.Now().Add(24 * time.Hour),
	}

	// 生成授权码
	authCode := generateAuthCode()

	// 重定向回客户端
	redirectURL := fmt.Sprintf("%s?code=%s&state=%s", redirectURI, authCode, state)
	ctx.Redirect(redirectURL, fasthttp.StatusFound)
}

// Token端点 - 颁发访问令牌
func tokenHandler(ctx *fasthttp.RequestCtx) {
	grantType := string(ctx.FormValue("grant_type"))
	//code := string(ctx.FormValue("code"))
	clientID := string(ctx.FormValue("client_id"))
	clientSecret := string(ctx.FormValue("client_secret"))

	if grantType != "authorization_code" {
		ctx.Error("Unsupported grant type", fasthttp.StatusBadRequest)
		return
	}

	if clientID != clientID_demo || clientSecret != clientSecret_demo {
		ctx.Error("Invalid client credentials", fasthttp.StatusUnauthorized)
		return
	}

	// 验证授权码
	//for sessionID, session := range sessions {
	for _, session := range sessions {
		// 	if sessionID == code && time.Now().Before(session.Expiry) {
		// 生成访问令牌和ID令牌
		accessToken := generateAccessToken()
		idToken := generateIDToken(session)

		response := TokenResponse{
			AccessToken: accessToken,
			TokenType:   "Bearer",
			ExpiresIn:   3600,
			IDToken:     idToken,
		}

		ctx.SetContentType("application/json")
		json.NewEncoder(ctx).Encode(response)
		return
		// 	}
	}

	// ctx.Error("Invalid authorization code", fasthttp.StatusUnauthorized)
}

// UserInfo端点 - 返回用户信息
func userInfoHandler(ctx *fasthttp.RequestCtx) {
	authHeader := string(ctx.Request.Header.Peek("Authorization"))
	if !strings.HasPrefix(authHeader, "Bearer ") {
		ctx.Error("Missing access token", fasthttp.StatusUnauthorized)
		return
	}

	token := strings.TrimPrefix(authHeader, "Bearer ")

	// 验证令牌并返回用户信息
	for _, session := range sessions {
		if generateAccessToken() == token && time.Now().Before(session.Expiry) {
			response := UserInfoResponse{
				Sub:               session.UserID,
				Name:              session.Name,
				Email:             session.Email,
				PreferredUsername: session.Username,
			}

			ctx.SetContentType("application/json")
			json.NewEncoder(ctx).Encode(response)
			return
		}
	}

	ctx.Error("Invalid access token", fasthttp.StatusUnauthorized)
}

// 根路径处理
func rootHandler(ctx *fasthttp.RequestCtx) {
	html_ := `
<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>本地OIDC服务</title>
    <style>
        body { font-family: Arial, sans-serif; max-width: 600px; margin: 50px auto; padding: 20px; }
        .endpoint { background: #f8f9fa; padding: 10px; margin: 10px 0; border-radius: 5px; }
        code { background: #e9ecef; padding: 2px 4px; border-radius: 3px; }
    </style>
</head>
<body>
    <h1>🎯 本地OIDC服务已启动</h1>
    
    <div class="endpoint">
        <h3>🔍 OIDC配置端点</h3>
        <p><a href="/.well-known/openid-configuration">%s/.well-known/openid-configuration</a></p>
    </div>
    
    <div class="endpoint">
        <h3>🔑 JWKS端点</h3>
        <p><a href="/.well-known/jwks.json">%s/.well-known/jwks.json</a></p>
    </div>
    
    <div class="endpoint">
        <h3>🔐 授权端点</h3>
        <p><a href="/authorize?client_id=%s&redirect_uri=%s&scope=openid profile email&state=test123">%s/authorize</a></p>
    </div>
    
    <div class="endpoint">
        <h3>🎫 Token端点</h3>
        <p><code>POST %s/token</code></p>
    </div>
    
    <div class="endpoint">
        <h3>👤 UserInfo端点</h3>
        <p><code>GET %s/userinfo</code></p>
    </div>
    
    <div class="endpoint">
        <h3>🧪 测试登录</h3>
        <p><a href="/login">%s/login</a></p>
    </div>
    
    <div class="endpoint">
        <h3>📋 测试用户</h3>
        <ul>
            <li><strong>admin</strong> / admin123</li>
            <li><strong>test</strong> / test123</li>
        </ul>
    </div>
</body>
</html>`
	ctx.SetContentType("text/html")
	fmt.Fprintf(ctx, html_, localIssuer, localIssuer, clientID, redirectURI, localIssuer, localIssuer, localIssuer, localIssuer, localIssuer)
}

// 工具函数
func generateSessionID() string {
	b := make([]byte, 32)
	rand.Read(b)
	return base64.URLEncoding.EncodeToString(b)
}

func generateAuthCode() string {
	b := make([]byte, 16)
	rand.Read(b)
	return base64.URLEncoding.EncodeToString(b)
}

func generateAccessToken() string {
	b := make([]byte, 32)
	rand.Read(b)
	return base64.URLEncoding.EncodeToString(b)
}

func initRSAKeys() error {
	var err error
	privateKey, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return fmt.Errorf("failed to generate RSA key: %v", err)
	}
	publicKey = &privateKey.PublicKey
	return nil
}

func generateIDToken(session *UserSession) string {
	claims := jwt.RegisteredClaims{
		Subject:   session.UserID,
		Issuer:    localIssuer,
		Audience:  jwt.ClaimStrings{clientID},
		ExpiresAt: jwt.NewNumericDate(time.Now().Add(1 * time.Hour)),
		IssuedAt:  jwt.NewNumericDate(time.Now()),
	}

	// 添加自定义声明
	customClaims := jwt.MapClaims{
		"name":               session.Name,
		"email":              session.Email,
		"preferred_username": session.Username,
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, customClaims)
	token.Claims = claims

	tokenString, err := token.SignedString(privateKey)
	if err != nil {
		return ""
	}

	return tokenString
}

func jwksHandler(ctx *fasthttp.RequestCtx) {
	// 生成JWK格式的公钥信息
	n := base64.RawURLEncoding.EncodeToString(publicKey.N.Bytes())
	e := base64.RawURLEncoding.EncodeToString([]byte{1, 0, 1}) // 65537

	jwk := map[string]any{
		"keys": []map[string]any{
			{
				"kty": "RSA",
				"kid": "demo-key-1",
				"use": "sig",
				"alg": "RS256",
				"n":   n,
				"e":   e,
			},
		},
	}

	ctx.SetContentType("application/json")
	json.NewEncoder(ctx).Encode(jwk)
}
